'''
Author: caoniu caoniu@jushuitan.com
Date: 2023-06-28 19:01:41
LastEditors: caoniu caoniu@jushuitan.com
LastEditTime: 2023-06-28 23:02:16
FilePath: /pdf_knowledge_base/pdf2vector.py
Description: 这是默认设置,请设置`customMade`, 打开koroFileHeader查看配置 进行设置: https://github.com/OBKoro1/koro1FileHeader/wiki/%E9%85%8D%E7%BD%AE
'''

import asyncio
from langchain.vectorstores import Pinecone
from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.document_loaders import PyPDFLoader
import pinecone
import openai
import sys
import os
from dotenv import load_dotenv

load_dotenv()

pdf_name = ''

if len(sys.argv) > 1:
    pdf_name = sys.argv[1]
else:
    while (len(pdf_name) == 0):
        pdf_name = input('\nEnter the full name of the PDF: ')
        if os.path.exists(f'./library/{pdf_name}') is False:
            print('\nFile not found.')
            pdf_name = ''

    if input(f'\nEnter \'yes\' if confirmed: ') != 'yes':
        exit()

res = openai.Embedding.create(
    api_key=os.environ["OPENAI_API_KEY"],
    input=[
        "Sample document text goes here",
        "there will be several phrases in each batch"
    ], engine='text-embedding-ada-002'
)
# extract embeddings to a list
embeds = [record['embedding'] for record in res['data']]
print(f'dimension: {len(embeds[0])}')

# initialize pinecone
pinecone.init(
    api_key=os.getenv("PINECONE_API_KEY"),  # find at app.pinecone.io
    environment=os.getenv("PINECONE_ENV")  # next to api key in console
)

index_name = 'knowledge-base'
# check if 'index_name' index already exists (only create index if not)
if index_name not in pinecone.list_indexes():
    print('Initializing index ...')
    pinecone.create_index(index_name, dimension=len(embeds[0]))
    print('Ready.')

print('Loading PDF ...')

loader = PyPDFLoader(f'./library/{pdf_name}')
documents = loader.load_and_split()
embeddings = OpenAIEmbeddings()
# loader = PyPDFDirectoryLoader("./library/")
# documents = loader.load_and_split()
# embeddings = OpenAIEmbeddings()

print('PDF loading complete')
print('Generating vector database ...')
Pinecone.from_documents(documents, embeddings, index_name=index_name)
print('Completed.')
